# memory_profiler.py

from collections import Counter
import functools
import weakref
from typing import Dict

import torch
from torch._subclasses import FakeTensorMode
from torch.utils._python_dispatch import TorchDispatchMode
from torch.utils._pytree import tree_map_only
from torch.utils.weak import WeakIdKeyDictionary


def tensor_storage_id(tensor):
    return tensor.storage()._cdata


class FakeTensorMemoryProfilerMode(TorchDispatchMode):
    def __init__(self):
        # counter of storage ids to live references
        self.storage_count: Dict[int, int] = Counter()
        # live fake tensors
        self.live_tensors = WeakIdKeyDictionary()
        # current memory usage
        self.memory_use = 0
        # maximum memory usage recorded
        self.max_memory = 0
        # list to track memory variations
        self.memory_varies = []
        # current function name being processed
        self.current_func_name = None
        # markers for important points in memory tracking
        self.markers = {}

    def __torch_dispatch__(self, func, types, args=(), kwargs=None):
        kwargs = kwargs if kwargs is not None else {}
        rs = func(*args, **kwargs)
        self.current_func_name = str(func)
        tree_map_only(torch._subclasses.FakeTensor, self.track_tensor_memory_use, rs)
        return rs

    def add_marker(self, marker_name):
        marker_val = len(self.memory_varies)
        self.markers[marker_name] = marker_val

    def track_tensor_memory_use(self, tensor):
        if tensor in self.live_tensors:
            return

        self.live_tensors[tensor] = True
        nbytes = tensor.untyped_storage().nbytes()
        storage_id = tensor_storage_id(tensor)

        if storage_id not in self.storage_count:
            self.change_memory(nbytes)

        self.storage_count[storage_id] += 1

        weakref.finalize(tensor, functools.partial(self.tensor_cleanup, storage_id, nbytes))

    def tensor_cleanup(self, storage_id, nbytes):
        self.storage_count[storage_id] -= 1
        if self.storage_count[storage_id] == 0:
            del self.storage_count[storage_id]
            self.change_memory(-nbytes)

    def change_memory(self, delta):
        self.memory_use += delta
        self.max_memory = max(self.memory_use, self.max_memory)
        self.memory_varies.append((self.current_func_name, self.memory_use))

    def current_memory_max(self):
        return self.max_memory

    def current_memory(self):
        return self.memory_use

    def draw_varies(self):
        func_names, memory_varies = list(zip(*self.memory_varies))
        for i in range(len(func_names)):
            print(i, func_names[i], memory_varies[i])
        x_ids = list(range(len(func_names)))
        import matplotlib.pyplot as plt
        plt.plot(x_ids, memory_varies, label="memory use")
        min_val = min(memory_varies)
        max_val = max(memory_varies)
        for marker_name, marker in self.markers.items():
            plt.plot([marker, marker], [min_val, max_val], lw=1, label=marker_name)
        plt.legend()

        plt.savefig("memory_varies.png", dpi=500)
        plt.show()
